import platform
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm
import os
from PIL import Image
import cv2
from pathlib import Path


def set_chinese_font():
    """根据操作系统设置中文字体"""
    system = platform.system()

    if system == "Windows":
        # Windows系统
        plt.rcParams["font.sans-serif"] = ["SimHei", "Microsoft YaHei", "DejaVu Sans"]
    elif system == "Darwin":
        # Mac系统
        plt.rcParams["font.sans-serif"] = [
            "Arial Unicode MS",
            "Heiti TC",
            "DejaVu Sans",
        ]
    else:
        # Linux系统
        plt.rcParams["font.sans-serif"] = ["WenQuanYi Micro Hei", "DejaVu Sans"]

    # 解决负号显示问题
    plt.rcParams["axes.unicode_minus"] = False
    print(f"已设置中文字体支持（{system}系统）")


def visualize_samples(
    target_base_dir=r"D:\desktop\class\DataSet\cats_vs_dogs_labeled", num_samples=10
):
    """可视化一些样本图片"""
    print("\n可视化样本图片...")

    # 获取一些猫和狗的图片路径
    train_cats_dir = os.path.join(target_base_dir, "train", "cats")
    train_dogs_dir = os.path.join(target_base_dir, "train", "dogs")

    cat_images = [
        os.path.join(train_cats_dir, f)
        for f in os.listdir(train_cats_dir)[: num_samples // 2]
    ]
    dog_images = [
        os.path.join(train_dogs_dir, f)
        for f in os.listdir(train_dogs_dir)[: num_samples // 2]
    ]

    # 创建图表
    fig, axes = plt.subplots(2, num_samples // 2, figsize=(15, 6))

    # 显示猫图片
    for i, img_path in enumerate(cat_images):
        img = Image.open(img_path)
        axes[0, i].imshow(img)
        axes[0, i].set_title("Cat")
        axes[0, i].axis("off")

    # 显示狗图片
    for i, img_path in enumerate(dog_images):
        img = Image.open(img_path)
        axes[1, i].imshow(img)
        axes[1, i].set_title("Dog")
        axes[1, i].axis("off")

    plt.tight_layout()
    plt.show()


def denormalize(tensor, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]):
    """
    反标准化,将Tensor转换回可显示的图片
    用来显式转换后转换前图像对比
    """
    tensor = tensor.clone()  # 不修改原Tensor
    for t, m, s in zip(tensor, mean, std):
        t.mul_(s).add_(m)  # 逆变换: (tensor * std) + mean
    return tensor


def show_transformed_images(num_images=8, train_dataset=None, transform=None):
    # 获取样本
    indices = np.random.choice(len(train_dataset), num_images, replace=False)

    fig, axes = plt.subplots(2, num_images, figsize=(20, 6))
    if num_images == 1:
        axes = axes.reshape(2, 1)

    for i, idx in enumerate(indices):
        # 获取原始图片路径
        image_path, label = train_dataset.samples[idx]

        # 显示原始图片
        original_img = Image.open(image_path).convert("L")
        axes[0, i].imshow(original_img, cmap="gray")
        axes[0, i].set_title(f"原始\n{train_dataset.classes[label]}")
        axes[0, i].axis("off")

        # 显示变换后图片
        transformed_img, _ = train_dataset[idx]

        # 直接显示，不进行反标准化（因为标准化不影响灰度显示）
        img_display = transformed_img.squeeze().numpy()  # [1, H, W] -> [H, W]

        axes[1, i].imshow(img_display, cmap="gray")
        axes[1, i].set_title(f"变换后\n{train_dataset.classes[label]}")
        axes[1, i].axis("off")

    plt.tight_layout()
    plt.show()


def convert_color_to_grayscale(source_dir, target_dir):
    """
    将源目录中的所有彩色图片转换为灰度图并保存到目标目录

    Args:
        source_dir: 源图片目录路径
        target_dir: 目标保存目录路径
    """
    # 创建目标目录（如果不存在）
    Path(target_dir).mkdir(parents=True, exist_ok=True)

    # 支持的图片格式
    supported_formats = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif"}

    # 遍历源目录中的所有文件
    for filename in os.listdir(source_dir):
        file_path = os.path.join(source_dir, filename)

        # 检查是否为图片文件
        if (
            os.path.isfile(file_path)
            and Path(filename).suffix.lower() in supported_formats
        ):
            try:
                # 读取彩色图片
                color_image = cv2.imread(file_path)

                if color_image is not None:
                    # 转换为灰度图
                    gray_image = cv2.cvtColor(color_image, cv2.COLOR_BGR2GRAY)

                    # 构建目标文件路径（保持原文件名）
                    target_path = os.path.join(target_dir, filename)

                    # 保存灰度图
                    cv2.imwrite(target_path, gray_image)
                    print(f"成功转换: {filename}")
                else:
                    print(f"无法读取图片: {filename}")

            except Exception as e:
                print(f"处理图片 {filename} 时出错: {str(e)}")


def crop_to_square_center(source_dir, target_dir, size=300):
    """
    将图片中心裁剪为指定大小的正方形

    Args:
        source_dir: 源图片目录路径
        target_dir: 目标保存目录路径
        size: 正方形边长,默认300
    """
    # 创建目标目录
    Path(target_dir).mkdir(parents=True, exist_ok=True)

    # 支持的图片格式
    supported_formats = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif"}

    # 遍历源目录中的所有文件
    for filename in os.listdir(source_dir):
        file_path = os.path.join(source_dir, filename)

        # 检查是否为图片文件
        if (
            os.path.isfile(file_path)
            and Path(filename).suffix.lower() in supported_formats
        ):
            try:
                # 读取图片
                img = cv2.imread(file_path)

                if img is not None:
                    # 获取图片尺寸
                    height, width = img.shape[:2]

                    # 计算裁剪区域
                    if width > height:
                        # 宽大于高，从宽度中心裁剪
                        start_x = (width - height) // 2
                        start_y = 0
                        crop_size = height
                    else:
                        # 高大于宽，从高度中心裁剪
                        start_x = 0
                        start_y = (height - width) // 2
                        crop_size = width

                    # 裁剪正方形区域
                    cropped_img = img[
                        start_y : start_y + crop_size, start_x : start_x + crop_size
                    ]

                    # 调整大小为300x300
                    resized_img = cv2.resize(cropped_img, (size, size))

                    # 构建目标文件路径
                    target_path = os.path.join(target_dir, filename)

                    # 保存图片
                    cv2.imwrite(target_path, resized_img)
                    print(f"成功裁剪: {filename} ({width}x{height} -> {size}x{size})")
                else:
                    print(f"无法读取图片: {filename}")

            except Exception as e:
                print(f"处理图片 {filename} 时出错: {str(e)}")


def force_convert_to_grayscale(root_dir):
    """
    强制将所有图片转换为真正的单通道
    """
    valid_extensions = (".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif")

    for root, dirs, files in os.walk(root_dir):
        for file in tqdm(files, desc=f"处理 {root}"):
            if file.lower().endswith(valid_extensions):
                file_path = os.path.join(root, file)

                try:
                    # 使用PIL打开并转换
                    with Image.open(file_path) as img:
                        # 强制转换为'L'模式（8位像素，黑白）
                        gray_img = img.convert("L")

                        # 保存时确保是单通道
                        # 对于PIL，保存为PNG时会保持单通道
                        gray_img.save(file_path)
                        print(f"已强制转换: {file_path}")

                except Exception as e:
                    print(f"错误: 处理 {file_path} 时出错 - {e}")
